{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b171b90d-c783-451a-b393-7e7b43a155d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import cv2\n",
    "import numpy as np\n",
    "import ctypes\n",
    "from ctypes import c_bool, c_double, c_float, c_int, c_ubyte, c_uint64, c_void_p, POINTER, Structure\n",
    "from gym import Env\n",
    "from gym.spaces import Box, Discrete\n",
    "from stable_baselines3 import DQN\n",
    "from stable_baselines3.common.monitor import Monitor\n",
    "from stable_baselines3.common.vec_env import DummyVecEnv\n",
    "from stable_baselines3.common.callbacks import BaseCallback\n",
    "import platform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b0c5f497-03e1-432e-a7fd-75216e02b921",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载DLL\n",
    "class Vector2(ctypes.Structure):\n",
    "    _fields_ = [\n",
    "        (\"x\", ctypes.c_double),\n",
    "        (\"y\", ctypes.c_double),\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bae8dbbe-3f0d-4c33-858f-4bda9820dbc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DinoInfo(ctypes.Structure):\n",
    "    _fields_ = [\n",
    "        (\"position\", Vector2),\n",
    "        (\"velocity\", Vector2),\n",
    "        (\"collision_box_lu\", Vector2),\n",
    "        (\"collision_box_rd\", Vector2),\n",
    "        (\"state\", ctypes.c_int),\n",
    "        (\"health\", ctypes.c_int64),\n",
    "        (\"energy\", ctypes.c_int64),\n",
    "    ]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f7bfe968-3414-478c-89e3-f11064563fee",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FireballInfo(ctypes.Structure):\n",
    "    _fields_ = [\n",
    "        (\"position\", Vector2),\n",
    "        (\"velocity\", Vector2),\n",
    "        (\"collision_box_lu\", Vector2),\n",
    "        (\"collision_box_rd\", Vector2),\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "793b9bd9-8416-493d-bbde-7383a98c0f8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PterosaurInfo(ctypes.Structure):\n",
    "    _fields_ = [\n",
    "        (\"position\", Vector2),\n",
    "        (\"velocity\", Vector2),\n",
    "        (\"collision_box_lu\", Vector2),\n",
    "        (\"collision_box_rd\", Vector2),\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "45244a77-cd97-40ca-bf8b-9acadd5bdf5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CactusInfo(ctypes.Structure):\n",
    "    _fields_ = [\n",
    "        (\"position\", Vector2),\n",
    "        (\"velocity\", Vector2),\n",
    "        (\"collision_box_lu\", Vector2),\n",
    "        (\"collision_box_rd\", Vector2),\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ea6bb491-7302-4892-8b78-74cc8432dc6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MeteoriteInfo(ctypes.Structure):\n",
    "    _fields_ = [\n",
    "        (\"position\", Vector2),\n",
    "        (\"velocity\", Vector2),\n",
    "        (\"collision_box_lu\", Vector2),\n",
    "        (\"collision_box_rd\", Vector2),\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "edd6bdf3-b03e-44b6-9236-beb94a8300b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GameInfo(ctypes.Structure):\n",
    "    _fields_ = [\n",
    "        (\"dinosaur\", ctypes.POINTER(DinoInfo)),\n",
    "        (\"fireballs\", ctypes.POINTER(FireballInfo)),\n",
    "        (\"fireball_count\", ctypes.c_uint64),\n",
    "        (\"pterosaurs\", ctypes.POINTER(PterosaurInfo)),\n",
    "        (\"pterosaur_count\", ctypes.c_uint64),\n",
    "        (\"cactuses\", ctypes.POINTER(CactusInfo)),\n",
    "        (\"cactus_count\", ctypes.c_uint64),\n",
    "        (\"meteorites\", ctypes.POINTER(MeteoriteInfo)),\n",
    "        (\"meteorite_count\", ctypes.c_uint64),\n",
    "        (\"flag_fail\", ctypes.c_bool),\n",
    "        (\"flag_pause\", ctypes.c_bool),\n",
    "        (\"score\", ctypes.c_uint64)\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b40f450d-d491-4386-84bc-5f44524cf364",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MiniDinoGame:\n",
    "    def __init__(self, to_render=False, to_extract_game_info=True, lib_path=None):\n",
    "        system_type = platform.system()\n",
    "        if lib_path is None:\n",
    "            if system_type == \"Windows\":\n",
    "                libfile_name = \"mini_dino.dll\"\n",
    "            elif system_type == \"Linux\":\n",
    "                libfile_name = \"libmini_dino.so\"\n",
    "            else:\n",
    "                raise RuntimeError(\"Unsupported os.\")\n",
    "\n",
    "        lib_path = \"D:/zhulibo/university/大三/智能系统控制/chrome-dinosaur/PyMini-Dino/mini_dino.dll\"\n",
    "        self.lib = ctypes.cdll.LoadLibrary(lib_path)\n",
    "\n",
    "        self.lib.create_game.argtypes = [c_bool, c_bool]\n",
    "        self.lib.create_game.restype = c_void_p\n",
    "\n",
    "        self.lib.initialize.argtypes = [c_void_p]\n",
    "        self.lib.initialize.restype = None\n",
    "\n",
    "        self.lib.run.argtypes = [c_void_p]\n",
    "        self.lib.run.restype = None\n",
    "\n",
    "        self.lib.step.argtypes = [c_void_p, c_void_p, c_double]\n",
    "        self.lib.step.restype = c_bool\n",
    "\n",
    "        self.lib.get_game_info_ptr.argtypes = [c_void_p]\n",
    "        self.lib.get_game_info_ptr.restype = ctypes.POINTER(GameInfo)\n",
    "\n",
    "        self.lib.get_buffer_ptr.argtypes = [c_void_p]\n",
    "        self.lib.get_buffer_ptr.restype = ctypes.POINTER(c_ubyte)\n",
    "\n",
    "        self.lib.reset.argtypes = [c_void_p]\n",
    "        self.lib.reset.restype = None\n",
    "\n",
    "        self.lib.clean_up.argtypes = [c_void_p]\n",
    "        self.lib.clean_up.restype = None\n",
    "\n",
    "        self.game_ptr = self.lib.create_game(to_render, to_extract_game_info)\n",
    "        if not self.game_ptr:\n",
    "            raise RuntimeError(\"Failed to init_game, got null pointer!\")\n",
    "\n",
    "        self.is_initialized = False\n",
    "\n",
    "    def initialize(self):\n",
    "        if not self.is_initialized:\n",
    "            self.lib.initialize(self.game_ptr)\n",
    "            self.is_initialized = True\n",
    "\n",
    "    def run(self):\n",
    "        if not self.is_initialized:\n",
    "            self.initialize()\n",
    "\n",
    "        # Check game status, break if paused\n",
    "        while not self.lib.step(self.game_ptr, None, 0.0167):\n",
    "            time.sleep(0.01)  # Sleep to avoid 100% CPU usage\n",
    "\n",
    "    def step(self, keyboard_ptr, delta_time: float):\n",
    "        # Check if the game failed or paused and handle accordingly\n",
    "        game_info = self.get_game_info()\n",
    "        if game_info['flag_fail']:\n",
    "            print(\"Game Over!\")\n",
    "            return False  # Stop game\n",
    "        elif game_info['flag_pause']:\n",
    "            print(\"Game Paused!\")\n",
    "            return False  # Game paused\n",
    "\n",
    "        return self.lib.step(self.game_ptr, keyboard_ptr, delta_time)\n",
    "\n",
    "    def reset(self):\n",
    "        self.lib.reset(self.game_ptr)\n",
    "\n",
    "    def get_game_info(self):\n",
    "        info_ptr = self.lib.get_game_info_ptr(self.game_ptr)\n",
    "        if not info_ptr:\n",
    "            print(\"Failed to get game info\")\n",
    "            return None\n",
    "        return info_ptr.contents\n",
    "\n",
    "    def get_game_frame(self):\n",
    "        \"\"\"\n",
    "        The returned image is stored in BGRA format,\n",
    "        which is common in OpenCV. Here we will ignore the alpha channel,\n",
    "        resize the image to a smaller size, and return it as RGB.\n",
    "        \"\"\"\n",
    "        buf_ptr = self.lib.get_buffer_ptr(self.game_ptr)\n",
    "        width, height = 1000, 600\n",
    "        buffer = np.ctypeslib.as_array(buf_ptr, shape=(width * height * 4,))  # 4 channels (BGRA)\n",
    "        image_data = buffer.reshape((height, width, 4))  # BGRA format\n",
    "        image_bgra = image_data.astype(np.uint8)\n",
    "    \n",
    "        # Ignore the alpha channel (last channel), keep only RGB\n",
    "        image_rgb = cv2.cvtColor(image_bgra, cv2.COLOR_BGRA2RGB)\n",
    "    \n",
    "        # Resize the image to 224x224 (or smaller)\n",
    "        resized_image = cv2.resize(image_rgb, (224, 224))\n",
    "    \n",
    "        return resized_image  # Now it's (224, 224, 3)\n",
    "\n",
    "\n",
    "    def create_keyboard_state(self, keyboard_dict):\n",
    "        keyboard_state = (ctypes.c_ubyte * 108)()\n",
    "        for key, pressed in keyboard_dict.items():\n",
    "            if 0 <= key < 108:\n",
    "                keyboard_state[key] = 1 if pressed else 0\n",
    "        return ctypes.cast(keyboard_state, ctypes.POINTER(ctypes.c_ubyte))\n",
    "\n",
    "    def close(self):\n",
    "        self.lib.clean_up(self.game_ptr)\n",
    "\n",
    "    def __del__(self):\n",
    "        self.close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "aa4b1a4f-860f-4ea4-b1fe-8705f4d19023",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DinoGameEnv(Env):\n",
    "    def __init__(self, to_render=False, to_extract_game_info=True, lib_path=None):\n",
    "        super().__init__()\n",
    "        self.game = MiniDinoGame(to_render=to_render, to_extract_game_info=to_extract_game_info, lib_path=lib_path)\n",
    "        self.game.initialize()\n",
    "\n",
    "        self.observation_space = Box(low=0, high=255, shape=(600, 1000, 3), dtype=np.uint8)\n",
    "        self.action_space = Discrete(5)  # w, a, s, d, j (5个操作)\n",
    "\n",
    "        self.last_score = 0\n",
    "\n",
    "    def reset(self):\n",
    "        self.game.reset()\n",
    "        time.sleep(1)\n",
    "        return self.get_observation()\n",
    "\n",
    "    def step(self, action):\n",
    "        keyboard_state = self.create_keyboard_state(action)\n",
    "        self.game.step(keyboard_state, 0.0167)\n",
    "\n",
    "        game_info = self.game.get_game_info()\n",
    "        done = game_info[\"flag_fail\"]\n",
    "        reward = self.calculate_reward(game_info)\n",
    "        observation = self.get_observation()\n",
    "\n",
    "        return observation, reward, done, {}\n",
    "\n",
    "    def render(self):\n",
    "        frame = self.game.get_game_frame()\n",
    "        cv2.imshow('Dino Game', frame)\n",
    "        if cv2.waitKey(1) & 0xFF == ord('q'):\n",
    "            self.close()\n",
    "\n",
    "    def close(self):\n",
    "        cv2.destroyAllWindows()\n",
    "\n",
    "    def get_observation(self):\n",
    "        return self.game.get_game_frame()\n",
    "\n",
    "    def create_keyboard_state(self, action):\n",
    "        keyboard_dict = {\n",
    "            0: {19: True},  # w\n",
    "            1: {52: True},  # a\n",
    "            2: {51: True},  # s\n",
    "            3: {53: True},  # d\n",
    "            4: {19: True}   # j\n",
    "        }\n",
    "        state = {key: False for key in range(108)}\n",
    "        if action in keyboard_dict:\n",
    "            for key, pressed in keyboard_dict[action].items():\n",
    "                state[key] = pressed\n",
    "        return self.game.create_keyboard_state(state)\n",
    "\n",
    "    def calculate_reward(self, game_info):\n",
    "        reward = 0\n",
    "        dino = game_info[\"dinosaur\"]\n",
    "\n",
    "        if dino:\n",
    "            health_reward = dino[\"health\"] / 100\n",
    "            energy_reward = dino[\"energy\"] / 100\n",
    "            score_reward = (game_info[\"score\"] - self.last_score) / 10\n",
    "\n",
    "            reward = health_reward + energy_reward + score_reward\n",
    "\n",
    "        if game_info[\"flag_fail\"]:\n",
    "            reward -= 1\n",
    "\n",
    "        self.last_score = game_info[\"score\"]\n",
    "        return reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7fc08ad7-98b1-41b2-beb3-ab69c5d62830",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrainAndLoggingCallback(BaseCallback):\n",
    "    def __init__(self, check_freq, save_path, verbose=1):\n",
    "        super(TrainAndLoggingCallback, self).__init__(verbose)\n",
    "        self.check_freq = check_freq\n",
    "        self.save_path = save_path\n",
    "\n",
    "    def _init_callback(self):\n",
    "        if self.save_path is not None:\n",
    "            os.makedirs(self.save_path, exist_ok=True)\n",
    "\n",
    "    def _on_step(self):\n",
    "        if self.n_calls % self.check_freq == 0:\n",
    "            model_path = os.path.join(self.save_path, 'best_model_{}'.format(self.n_calls))\n",
    "            self.model.save(model_path)\n",
    "\n",
    "        return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3a14f3f9-5f14-4f6b-bc1d-22eff39b563d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 环境设置\n",
    "CHECKPOINT_DIR = './train/'\n",
    "LOG_DIR = './logs/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "69df95ff-7e6e-4417-ad64-51f45724123b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "D:\\zhulibo\\software\\anaconda3\\envs\\dino\\lib\\site-packages\\stable_baselines3\\common\\vec_env\\patch_gym.py:49: UserWarning: You provided an OpenAI Gym environment. We strongly recommend transitioning to Gymnasium environments. Stable-Baselines3 is automatically wrapping your environments in a compatibility layer, which could potentially cause issues.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "env = DummyVecEnv([lambda: DinoGameEnv()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "97bdd9ea-6d6f-424b-a99c-eb935423314b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using cpu device\n",
      "Wrapping the env in a VecTransposeImage.\n"
     ]
    }
   ],
   "source": [
    "# 创建并训练模型\n",
    "model = DQN('CnnPolicy', env, tensorboard_log=LOG_DIR, verbose=1, buffer_size=1000, learning_starts=200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44e2a7f7-4554-4f9e-9a38-0c7f116f678f",
   "metadata": {},
   "outputs": [],
   "source": [
    "callback = TrainAndLoggingCallback(check_freq=1000, save_path=CHECKPOINT_DIR)\n",
    "model.learn(total_timesteps=10000, callback=callback)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdd6ed93-444b-4c29-9617-39f6ae4fb5c6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
